// test_concurrent.c
#include "test_framework.h"
#include <linux/kthread.h>
#include <linux/completion.h>
#include "../hisi_l0_mem_pool.h"
#include "test_common.h"
#include "test_registration.h"

struct thread_test_data {
    struct task_struct *task;
    struct completion completion;
    int thread_id;
    int iterations;
    atomic_t *error_count;
};

static int allocation_thread(void *data)
{
    struct thread_test_data *tdata = (struct thread_test_data *)data;
    int i;
    void *ptr;
    
    for (i = 0; i < tdata->iterations; i++) {
        int size = 32 + (i % 10) * 32; // Varying sizes
        
        ptr = l0_kmalloc(size, 0);
        if (!ptr) {
            atomic_inc(tdata->error_count);
            break;
        }
        memset(ptr, tdata->thread_id & 0xFF, size);
        l0_kfree(ptr);
    }
    
    complete(&tdata->completion);
    return 0;
}

/**
 * test_concurrent_allocations - Tests concurrent allocation operations
 * 
 * This test verifies:
 * 1. Thread safety of allocation operations
 * 2. No race conditions in multi-threaded scenarios
 * 3. System stability under concurrent allocation load
 * 4. Proper cleanup in concurrent environments
 * 
 * Tests multi-threaded allocation safety
 */
static int test_concurrent_allocations(struct seq_file *m)
{
#define NUM_THREADS 4
#define ITERATIONS 10
    
    struct thread_test_data threads[NUM_THREADS];
    atomic_t error_count;
    int i;
    int passed = 0, total = 0;
    
    atomic_set(&error_count, 0);
    
    // Create multiple threads performing allocations
    for (i = 0; i < NUM_THREADS; i++) {
        total++;
        init_completion(&threads[i].completion);
        threads[i].thread_id = i;
        threads[i].iterations = ITERATIONS;
        threads[i].error_count = &error_count;
        
        threads[i].task = kthread_run(allocation_thread, &threads[i], 
                                     "l0_test_thread_%d", i);
        if (IS_ERR(threads[i].task)) {
            TEST_FAIL(m, "Failed to create thread %d", i);
            continue;
        }
        passed++;
    }
    
    // Wait for all threads to complete
    for (i = 0; i < NUM_THREADS; i++) {
        if (!IS_ERR(threads[i].task)) {
            wait_for_completion(&threads[i].completion);
        }
    }
    
    // Check for errors
    if (atomic_read(&error_count) == 0) {
        TEST_PASS(m, "Concurrent allocations successful with %d threads", NUM_THREADS);
    } else {
        TEST_FAIL(m, "Concurrent allocations failed with %d errors", atomic_read(&error_count));
    }
    
    TEST_END(m, "Concurrent allocation tests: %d/%d passed", passed, total);
    return (passed == total && atomic_read(&error_count) == 0) ? 0 : -1;
    
#undef NUM_THREADS
#undef ITERATIONS
}

struct complex_thread_data {
    struct task_struct *task;
    struct completion completion;
    int thread_id;
    int iterations;
    atomic_t *error_count;
    int test_node_id;
};

static int complex_allocation_thread(void *data)
{
    struct complex_thread_data *tdata = (struct complex_thread_data *)data;
    int i;
    void *ptrs[10];
    
    for (i = 0; i < tdata->iterations; i++) {
        int size, j;
        int num_allocs = (i % 5) + 1; // 1-5 allocations per iteration
        
        // Allocate multiple blocks of varying sizes
        for (j = 0; j < num_allocs; j++) {
            size = 32 + (i * j % 32) * 32; // Varying sizes
            ptrs[j] = l0_kmalloc(size, tdata->test_node_id);
            if (!ptrs[j]) {
                atomic_inc(tdata->error_count);
                // Free any previously allocated blocks in this iteration
                while (--j >= 0) {
                    if (ptrs[j]) l0_kfree(ptrs[j]);
                }
                break;
            }
            memset(ptrs[j], (tdata->thread_id + i + j) & 0xFF, size);
        }
        
        // Free all allocated blocks
        for (j = 0; j < num_allocs; j++) {
            if (ptrs[j]) l0_kfree(ptrs[j]);
        }
    }
    
    complete(&tdata->completion);
    return 0;
}

/**
 * test_complex_concurrent_allocations - Tests complex concurrent operations
 * 
 * This test verifies:
 * 1. Complex multi-threaded allocation/deallocation patterns
 * 2. Thread safety with varying allocation sizes
 * 3. System stability under heavy concurrent load
 * 4. Proper resource management in concurrent scenarios
 * 
 * Tests advanced concurrent allocation scenarios
 */
static int test_complex_concurrent_allocations(struct seq_file *m)
{
#define NUM_THREADS 8
#define ITERATIONS 50
    
    struct complex_thread_data threads[NUM_THREADS];
    atomic_t error_count;
    int i;
    int passed = 0, total = 0;
    
    atomic_set(&error_count, 0);
    
    // Create multiple threads performing complex allocations
    for (i = 0; i < NUM_THREADS; i++) {
        total++;
        init_completion(&threads[i].completion);
        threads[i].thread_id = i;
        threads[i].iterations = ITERATIONS;
        threads[i].error_count = &error_count;
        threads[i].test_node_id = get_test_context()->test_node_id;
        
        threads[i].task = kthread_run(complex_allocation_thread, &threads[i], 
                                     "l0_complex_thread_%d", i);
        if (IS_ERR(threads[i].task)) {
            TEST_FAIL(m, "Failed to create complex thread %d", i);
            continue;
        }
        passed++;
    }
    
    // Wait for all threads to complete
    for (i = 0; i < NUM_THREADS; i++) {
        if (!IS_ERR(threads[i].task)) {
            wait_for_completion(&threads[i].completion);
        }
    }
    
    // Check for errors
    if (atomic_read(&error_count) == 0) {
        TEST_PASS(m, "Complex concurrent allocations successful with %d threads", NUM_THREADS);
    } else {
        TEST_FAIL(m, "Complex concurrent allocations failed with %d errors", atomic_read(&error_count));
    }
    
    TEST_END(m, "Complex concurrent allocation tests: %d/%d passed", passed, total);
    return (passed == total && atomic_read(&error_count) == 0) ? 0 : -1;
    
#undef NUM_THREADS
#undef ITERATIONS
}

// Export test cases
static struct test_case test_concurrent_cases[] = {
    {"Concurrent Allocations", test_concurrent_allocations},
    {"Complex Concurrent", test_complex_concurrent_allocations}
};

DECLARE_TEST_MODULE(concurrent, test_concurrent_cases)